{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import wandb\n",
    "import numpy as np\n",
    "import pandas as pd \n",
    "import matplotlib.pyplot as plt\n",
    "from tqdm import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def _get_summary(config):\n",
    "    summary = {\n",
    "        \"initial_lr\": config[\"init_lr\"],\n",
    "        \"mask_ratio\": config[\"mask_ratio\"],\n",
    "        \"d_model\": config[\"d_model\"],\n",
    "        \"train_batch_size\": config[\"train_batch_size\"],\n",
    "        \"patch_length\": config[\"patch_len\"],\n",
    "        \"torch_dtype\": config[\"torch_dtype\"],\n",
    "        \"train_ratio\": config[\"train_ratio\"],\n",
    "        \"val_ratio\": config[\"val_ratio\"],\n",
    "        \"test_ratio\": config[\"test_ratio\"],\n",
    "        \"random_seed\": config[\"random_seed\"],\n",
    "        \"model_name\": config[\"model_name\"],\n",
    "    }\n",
    "\n",
    "    if \"transformer_type\" in config:\n",
    "        summary[\"transformer_type\"] = config[\"transformer_type\"]\n",
    "    else:\n",
    "        summary[\"transformer_type\"] = \"None\"\n",
    "  \n",
    "    if \"randomly_initialize_backbone\" in config:\n",
    "        summary[\"randomly_initialize_backbone\"] = config[\"randomly_initialize_backbone\"]\n",
    "    else:\n",
    "        summary[\"randomly_initialize_backbone\"] = \"None\"\n",
    "    \n",
    "    if \"transformer_backbone\" in config:\n",
    "        summary[\"transformer_backbone\"] = config[\"transformer_backbone\"]\n",
    "    else:\n",
    "        summary[\"transformer_backbone\"] = \"None\"\n",
    "\n",
    "    if \"dataset_names\" in config and isinstance(config[\"dataset_names\"], str):\n",
    "        summary[\"dataset_names\"] = config[\"dataset_names\"]\n",
    "    else:\n",
    "        summary[\"dataset_names\"] = \"Multiple datasets.\"\n",
    "    return summary\n",
    "\n",
    "def _get_metadata_summary(metadata):\n",
    "    meta = {\n",
    "        \"hostname\": metadata[\"host\"],\n",
    "        \"gpu_type\": metadata[\"gpu\"],\n",
    "        \"program\": metadata[\"program\"],\n",
    "        \"args\": metadata[\"args\"],\n",
    "        \"username\": metadata[\"username\"],\n",
    "        \"execution_start_time\": metadata[\"startedAt\"],\n",
    "        \"git_commit\": metadata['git']['commit'],\n",
    "    }\n",
    "    if \"codePath\" in metadata:\n",
    "        meta[\"code_path\"] = metadata[\"codePath\"]\n",
    "    else:\n",
    "        meta[\"code_path\"] = \"None\"\n",
    "    return meta"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "api = wandb.Api()\n",
    "\n",
    "runs = api.runs(\"timeseries-foundation-model/Time-series Foundation Model\")\n",
    "\n",
    "summary_list = []\n",
    "for run in tqdm(runs, total=len(runs)): \n",
    "    if \"dataset_names\" not in run.config: \n",
    "        continue\n",
    "    \n",
    "    try:\n",
    "        summary = _get_summary(run.config)\n",
    "        metadata = _get_metadata_summary(run.metadata)\n",
    "    except:\n",
    "        print(f\"Failed to get summary for {run.name}\")\n",
    "    summary = {**summary, **metadata}\n",
    "    \n",
    "    summary[\"run_name\"] = run.name \n",
    "    summary[\"run_id\"] = run.id\n",
    "    summary[\"notes\"] = run.notes\n",
    "    summary_list.append(summary)\n",
    "\n",
    "columns = ['run_name', 'transformer_backbone', 'd_model',\n",
    "           'model_name', 'transformer_type',  'randomly_initialize_backbone', \n",
    "           'patch_length', 'random_seed', 'initial_lr', \n",
    "           'mask_ratio', 'train_batch_size', 'torch_dtype', \n",
    "           'train_ratio', 'val_ratio', 'test_ratio', \n",
    "           'hostname', 'git_commit', 'notes', 'gpu_type',\n",
    "           'program', 'code_path', 'args', 'username', \n",
    "           'execution_start_time', 'dataset_names']\n",
    "\n",
    "summary = pd.DataFrame(summary_list)\n",
    "summary.set_index(\"run_id\", inplace=True)\n",
    "summary = summary[columns]\n",
    "summary.to_csv(\"../../assets/data/wandb_runs_summary.csv\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "summary.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pretraining_runs = [\n",
    "    \"fearless-planet-52\", # Old runs\n",
    "    \"charmed-bee-241\", \"zesty-frost-243\", \"hopeful-water-240\", # New pre-training runs\n",
    "    \"Pre-training - small - patch 4\", \"Pre-training - small - patch 8, pos emb False\" # Ablation experiments\n",
    "]\n",
    "\n",
    "summary[summary[\"run_name\"].isin(pretraining_runs)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "### Pre-training runs\n",
    "small_run = api.run(path=os.path.join(run.entity, run.project, \"xxrzrfrd\")) # Small run\n",
    "base_run = api.run(path=os.path.join(run.entity, run.project, \"ligbv4fg\")) # Base run\n",
    "large_run = api.run(path=os.path.join(run.entity, run.project, \"0diyru56\")) # Large run\n",
    "\n",
    "# old_large_run = api.run(path=os.path.join(run.entity, run.project, \"new8a30j\")) # Old large run\n",
    "\n",
    "# Ablations\n",
    "small_patch4_run = api.run(path=os.path.join(run.entity, run.project, \"wf046knr\")) # Small run with patch_len = 4\n",
    "small_flant5_run = api.run(path=os.path.join(run.entity, run.project, \"hvyei2nf\")) # Small run with Flan T5 initialization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get step losses\n",
    "small_vals = pd.DataFrame(small_run.scan_history(keys=['_step', 'step_train_loss']))\n",
    "base_vals = pd.DataFrame(base_run.scan_history(keys=['_step', 'step_train_loss']))\n",
    "large_vals = pd.DataFrame(large_run.scan_history(keys=['_step', 'step_train_loss']))\n",
    "\n",
    "# old_large_vals = pd.DataFrame(old_large_run.scan_history(keys=['_step', 'step_train_loss']))\n",
    "\n",
    "small_patch4_vals = pd.DataFrame(small_patch4_run.scan_history(keys=['_step', 'step_train_loss']))\n",
    "small_flant5_vals = pd.DataFrame(small_flant5_run.scan_history(keys=['_step', 'step_train_loss']))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "span = 500\n",
    "\n",
    "small_losses = small_vals[\"step_train_loss\"].ewm(span=span).mean().to_numpy()\n",
    "base_losses = base_vals[\"step_train_loss\"].ewm(span=span).mean().to_numpy()\n",
    "large_losses = large_vals[\"step_train_loss\"].ewm(span=span).mean().to_numpy()\n",
    "\n",
    "# old_large_losses = old_large_vals[\"step_train_loss\"].ewm(span=span).mean().to_numpy()\n",
    "\n",
    "small_patch4_losses = small_patch4_vals[\"step_train_loss\"].ewm(span=span).mean().to_numpy()\n",
    "small_flant5_losses = small_flant5_vals[\"step_train_loss\"].ewm(span=span).mean().to_numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "billion = 1e9\n",
    "num_patches = 64*2048 # NOTE: x-axis is only approximate, since some patches may not contain any information. Also patches are repeated\n",
    "steps_in_one_epoch = num_patches*38374/billion\n",
    "\n",
    "plt.plot(num_patches*np.arange(1, len(small_losses)+1)/billion, \n",
    "         small_losses, c='red', label=\"  40M Small\")\n",
    "plt.plot(num_patches*np.arange(1, len(base_losses)+1)/billion, \n",
    "         base_losses, c='blue', label=\"125M Base\")\n",
    "plt.plot(num_patches*np.arange(1, len(large_losses)+1)/billion, \n",
    "         large_losses, c='green', label=\"385M Large\")\n",
    "\n",
    "# Epochs\n",
    "plt.axvline(x=steps_in_one_epoch, c='black', linestyle='--', linewidth=0.5)\n",
    "plt.axvline(x=2*steps_in_one_epoch, c='black', linestyle='--', linewidth=0.5)\n",
    "plt.axvline(x=3*steps_in_one_epoch, c='black', linestyle='--', linewidth=0.5)\n",
    "\n",
    "# Old run\n",
    "# plt.plot(num_patches*np.arange(1, len(old_large_losses)+1)/billion, \n",
    "#          old_large_losses, c='pink', label=\"385M Large (O)\")\n",
    "\n",
    "plt.grid(color='lightgray', linestyle='--', linewidth=0.5)  # Set grid color and style\n",
    "\n",
    "plt.ylim(0.06, 0.2)\n",
    "plt.xticks(fontsize=18)\n",
    "plt.yticks(fontsize=18)\n",
    "plt.xlabel(\"Processed patches (in Billions)\", fontsize=20)\n",
    "plt.ylabel(\"Training Loss\", fontsize=20)\n",
    "plt.xlim(0, 2*steps_in_one_epoch)\n",
    "plt.legend(fontsize=20)\n",
    "plt.savefig(\"../../assets/figures/pretraining/moment_family_training_losses.pdf\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(num_patches*np.arange(1, len(small_flant5_losses)+1)/billion, \n",
    "         small_flant5_losses, c='red', label=\"Small (Flan-T5)\")\n",
    "plt.plot(num_patches*np.arange(1, len(small_losses)+1)/billion, \n",
    "         small_losses, c='blue', label=\"Small (Random)\")\n",
    "\n",
    "# Epochs\n",
    "plt.axvline(x=steps_in_one_epoch, c='black', linestyle='--', linewidth=0.5)\n",
    "plt.axvline(x=2*steps_in_one_epoch, c='black', linestyle='--', linewidth=0.5)\n",
    "plt.axvline(x=3*steps_in_one_epoch, c='black', linestyle='--', linewidth=0.5)\n",
    "\n",
    "plt.ylim(0.08, 0.2)\n",
    "plt.xticks(fontsize=18)\n",
    "plt.yticks(fontsize=18)\n",
    "plt.grid(color='lightgray', linestyle='--', linewidth=0.5)\n",
    "plt.xlabel(\"Processed patches (in Billions)\", fontsize=20)\n",
    "plt.ylabel(\"Training Loss\", fontsize=20)\n",
    "plt.xlim(0, 2*steps_in_one_epoch)\n",
    "plt.legend(fontsize=20)\n",
    "plt.savefig(\"../../assets/figures/pretraining/random_vs_flant5_initialization.pdf\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(num_patches*np.arange(1, len(small_flant5_losses)+1)/billion, \n",
    "         small_flant5_losses, c='red', label=\"Small (Flan-T5)\")\n",
    "plt.plot(num_patches*np.arange(1, len(small_losses)+1)/billion, \n",
    "         small_losses, c='blue', label=\"Small (Random)\")\n",
    "\n",
    "# Epochs\n",
    "plt.axvline(x=steps_in_one_epoch, c='black', linestyle='--', linewidth=0.5)\n",
    "plt.axvline(x=2*steps_in_one_epoch, c='black', linestyle='--', linewidth=0.5)\n",
    "plt.axvline(x=3*steps_in_one_epoch, c='black', linestyle='--', linewidth=0.5)\n",
    "\n",
    "plt.xscale(\"log\")\n",
    "plt.xlim(1e-3, None)\n",
    "plt.ylim(0.06, 2)\n",
    "plt.xticks(fontsize=18)\n",
    "plt.yticks(fontsize=18)\n",
    "plt.grid(color='lightgray', linestyle='--', linewidth=0.5)\n",
    "plt.xlabel(\"Processed patches (in Billions)\", fontsize=20)\n",
    "plt.ylabel(\"Training Loss\", fontsize=20)\n",
    "plt.xlim(0, 2*steps_in_one_epoch)\n",
    "plt.legend(fontsize=20)\n",
    "plt.savefig(\"../../assets/figures/pretraining/random_vs_flant5_initialization_log.pdf\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(num_patches*np.arange(1, len(small_patch4_losses)+1)/billion, \n",
    "         small_patch4_losses, c='red', label=\"Small (Patch 4)\")\n",
    "plt.plot(num_patches*np.arange(1, len(small_losses)+1)/billion, \n",
    "         small_losses, c='blue', label=\"Small (Patch 8)\")\n",
    "\n",
    "# Epochs\n",
    "plt.axvline(x=steps_in_one_epoch, c='black', linestyle='--', linewidth=0.5)\n",
    "plt.axvline(x=2*steps_in_one_epoch, c='black', linestyle='--', linewidth=0.5)\n",
    "plt.axvline(x=3*steps_in_one_epoch, c='black', linestyle='--', linewidth=0.5)\n",
    "\n",
    "\n",
    "plt.ylim(0.07, 0.2)\n",
    "plt.xticks(fontsize=18)\n",
    "plt.yticks(fontsize=18)\n",
    "plt.grid(color='lightgray', linestyle='--', linewidth=0.5)\n",
    "plt.xlabel(\"Processed patches (in Billions)\", fontsize=20)\n",
    "plt.ylabel(\"Training Loss\", fontsize=20)\n",
    "plt.xlim(0, 2*steps_in_one_epoch)\n",
    "plt.legend(fontsize=20)\n",
    "plt.savefig(\"../../assets/figures/pretraining/patch_len_4_vs_8.pdf\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Radar Plot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.ticker as ticker"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "categories = [\n",
    "    'Imputation',\n",
    "    \"Long-horizon \\nForecasting\",\n",
    "    'Short-horizon \\nForecasting',\n",
    "    'Classification',\n",
    "    'Anomaly\\nDetection',\n",
    "]\n",
    "\n",
    "model_values = {\n",
    "    'MOMENT':     [\n",
    "        0.9929824561, \n",
    "        0.9958246347, \n",
    "        0.8849206349,\n",
    "        0.236097561, \n",
    "        0.9035087719],\n",
    "    'GPT4TS':   [\n",
    "        1, \n",
    "        0.9826026444, \n",
    "        0.4345238095,\n",
    "        0.3502439024, \n",
    "        0],\n",
    "    'TimesNet': [\n",
    "        0.9912280702, \n",
    "        0.9485038274, \n",
    "        0.4464285714,\n",
    "        0.1882926829, \n",
    "        0.798245614],\n",
    "}\n",
    "colors = [\n",
    "    'orangered',\n",
    "    'mediumblue',\n",
    "    'limegreen',\n",
    "]\n",
    "\n",
    "linestyles = ['solid', '--', ':']\n",
    "\n",
    "N = len(categories)\n",
    "angles = [n / float(N) * 2 * np.pi for n in range(N)]\n",
    "angles += angles[:1]\n",
    "\n",
    "fig, ax = plt.subplots(subplot_kw=dict(polar=True))\n",
    "ax.set_rlabel_position(0)\n",
    "plt.xticks(angles[:-1], [])\n",
    "plt.yticks([0.0, 0.33, 0.67, 1.0], [])\n",
    "ax.tick_params(axis='x', which='major', pad=16, labelsize=16)\n",
    "plt.ylim(0, 1.0)\n",
    "\n",
    "plt.text(0, 1.35, 'Imputation', fontsize=16, ha='center')\n",
    "plt.text(1, 1.25, \"Long-horizon Forecasting\", fontsize=16, ha='center')\n",
    "plt.text(2.7, 1.35, \"Classification\", fontsize=16, ha='center')\n",
    "plt.text(3.66, 1.37, \"Short-horizon\\nForecasting\", fontsize=16, ha='center')\n",
    "plt.text(5.25, 1.3, \"Anomaly Detection\", fontsize=16, ha='center')\n",
    "\n",
    "for model_name, values in model_values.items():\n",
    "    values_with_first = values + values[:1]  # close the circle\n",
    "    c = colors.pop(0)\n",
    "    linestyle = linestyles.pop(0)\n",
    "    ax.plot(angles, values_with_first, label=model_name, color=c, linewidth=1, linestyle=linestyle)\n",
    "    ax.fill(angles, values_with_first, alpha=0.5, color=c)\n",
    "\n",
    "plt.grid(color='lightgray', linestyle='--', linewidth=1)\n",
    "legend = plt.legend(fontsize=16, ncol=3, loc='lower center', bbox_to_anchor=(0.5, -0.25))\n",
    "# get the lines and texts inside legend box\n",
    "legend_lines = legend.get_lines()\n",
    "# bulk-set the properties of all lines and texts\n",
    "plt.setp(legend_lines, linewidth=3)\n",
    "plt.savefig(\"../../assets/figures/pretraining/model_comparison.pdf\", bbox_inches='tight')\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
